#!/usr/bin/env bash
# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
# Licensed under the CANN Open Software License Agreement Version 1.0.
# See LICENSE in the root of the software repository for full license details.

# =================================================================
# GELU Operator Test Script
# This script:
# 1. Sets up the Ascend environment
# 2. Generates test data
# 3. Builds and runs the GELU operator
# 4. Verifies the results
# =================================================================

set -eo pipefail

# --------------------------
# Configuration
# --------------------------
SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
LOG_DIR="${HOME}/ascend/log"
INPUT_DIR="${SCRIPT_DIR}/input"
OUTPUT_DIR="${SCRIPT_DIR}/output"
BUILD_DIR="${SCRIPT_DIR}/build"

# --------------------------
# Environment Setup
# --------------------------
setup_environment() {
    local ascend_path=""
    
    # Determine Ascend installation path
    if [[ -n "${ASCEND_INSTALL_PATH}" ]]; then
        ascend_path="${ASCEND_INSTALL_PATH}"
    elif [[ -n "${ASCEND_HOME_PATH}" ]]; then
        ascend_path="${ASCEND_HOME_PATH}"
    elif [[ -d "${HOME}/Ascend/ascend-toolkit/latest" ]]; then
        ascend_path="${HOME}/Ascend/ascend-toolkit/latest"
    else
        ascend_path="/usr/local/Ascend/ascend-toolkit/latest"
    fi

    # Verify Ascend installation
    if [[ ! -d "${ascend_path}" ]]; then
        echo "ERROR: Ascend Toolkit not found at ${ascend_path}" >&2
        return 1
    fi

    # Source environment settings
    if [[ -f "${ascend_path}/bin/setenv.bash" ]]; then
        source "${ascend_path}/bin/setenv.bash"
    else
        echo "WARNING: setenv.bash not found in ${ascend_path}/bin" >&2
    fi

    export DDK_PATH="${ascend_path}"
    export NPU_HOST_LIB="${ascend_path}/lib64"
    
    echo "INFO: Using Ascend Toolkit at ${ascend_path}"
}

# --------------------------
# Cleanup Directories
# --------------------------
clean_directories() {
    echo "INFO: Cleaning working directories..."
    
    # Create directories if they don't exist
    mkdir -p "${LOG_DIR}" "${INPUT_DIR}" "${OUTPUT_DIR}" "${BUILD_DIR}"
    
    # Clean log files
    find "${LOG_DIR}" -type f -delete
    
    # Clean input/output binaries
    rm -f "${INPUT_DIR}"/*.bin
    rm -f "${OUTPUT_DIR}"/*.bin
}

# --------------------------
# Generate Test Data
# --------------------------
generate_test_data() {
    echo "INFO: Generating test data..."
    
    if ! python3 "${SCRIPT_DIR}/gen_data.py"; then
        echo "ERROR: Failed to generate test data" >&2
        return 1
    fi
    
    echo "INFO: Test data generated successfully"
}

# --------------------------
# Build and Run
# --------------------------
build_and_run() {
    echo "INFO: Building GELU operator..."
    
    pushd "${BUILD_DIR}" > /dev/null
    
    # Configure and build
    if ! cmake .. && cmake --build . -j$(nproc); then
        echo "ERROR: Build failed" >&2
        popd > /dev/null
        return 1
    fi
    
    # Run the operator
    echo "INFO: Running GELU operator..."
    if ! ./execute_gelu_op; then
        echo "ERROR: Operator execution failed" >&2
        popd > /dev/null
        return 1
    fi
    
    popd > /dev/null
}

# --------------------------
# Verify Results
# --------------------------
verify_results() {
    echo "INFO: Verifying results..."
    
    local result
    result=$(python3 "${SCRIPT_DIR}/verify_result.py" \
        "${OUTPUT_DIR}/output_z.bin" \
        "${OUTPUT_DIR}/golden.bin")
    
    echo "RESULT: ${result}"
    
    if [[ "${result}" == "test pass" ]]; then
        echo ""
        echo "#####################################"
        echo "INFO: Precision test passed!"
        echo "#####################################"
        echo ""
        return 0
    else
        echo "ERROR: Precision test failed!" >&2
        return 1
    fi
}

# --------------------------
# Main Execution
# --------------------------
main() {
    # Setup environment
    if ! setup_environment; then
        exit 1
    fi
    
    # Clean working directories
    clean_directories
    
    # Generate test data
    if ! generate_test_data; then
        exit 1
    fi
    
    # Build and run
    if ! build_and_run; then
        exit 1
    fi
    
    # Verify results
    if ! verify_results; then
        exit 1
    fi
    
    echo "INFO: All steps completed successfully"
}

# Execute main function
main "$@"